// Copyright (c) 2024 IndustrialTBB Framework
// Licensed under MIT License

#ifndef ITBB_CORE_TBB_TASK_SCHEDULER_H_
#define ITBB_CORE_TBB_TASK_SCHEDULER_H_

#include <oneapi/tbb.h>
#include <oneapi/tbb/global_control.h>
#include <oneapi/tbb/task_arena.h>
#include <oneapi/tbb/task_scheduler_observer.h>

#include <atomic>
#include <chrono>
#include <functional>
#include <future>
#include <memory>

namespace itbb {
// TBB原生的任务优先级系统
enum class TaskPriority { kLow = 0, kNormal = 1, kHigh = 2, kCritical = 3 };

// TBB原生的任务统计信息
struct TbbTaskStats {
  std::atomic<uint64_t> total_tasks_submitted{0};
  std::atomic<uint64_t> total_tasks_completed{0};
  std::atomic<uint64_t> total_tasks_failed{0};
  std::atomic<uint64_t> total_execution_time_ns{0};
  std::atomic<uint64_t> max_execution_time_ns{0};
  std::atomic<uint64_t> min_execution_time_ns{UINT64_MAX};

  // 按优先级统计
  std::array<std::atomic<uint64_t>, 4> tasks_by_priority{};

  void RecordTaskExecution(TaskPriority priority,
                           std::chrono::nanoseconds duration) {
    total_tasks_completed.fetch_add(1, std::memory_order_relaxed);
    tasks_by_priority[static_cast<size_t>(priority)].fetch_add(
        1, std::memory_order_relaxed);

    uint64_t duration_ns = duration.count();
    total_execution_time_ns.fetch_add(duration_ns, std::memory_order_relaxed);

    // 更新最大执行时间
    uint64_t current_max =
        max_execution_time_ns.load(std::memory_order_acquire);
    while (duration_ns > current_max &&
           !max_execution_time_ns.compare_exchange_weak(
               current_max, duration_ns, std::memory_order_release)) {
      current_max = max_execution_time_ns.load(std::memory_order_acquire);
    }

    // 更新最小执行时间
    uint64_t current_min =
        min_execution_time_ns.load(std::memory_order_acquire);
    while (duration_ns < current_min &&
           !min_execution_time_ns.compare_exchange_weak(
               current_min, duration_ns, std::memory_order_release)) {
      current_min = min_execution_time_ns.load(std::memory_order_acquire);
    }
  }
};

// TBB原生的任务包装器
template <typename Callable>
class TbbTaskWrapper {
 public:
  TbbTaskWrapper(Callable&& callable, TaskPriority priority,
                 TbbTaskStats& stats)
      : callable_(std::forward<Callable>(callable)),
        priority_(priority),
        stats_(stats),
        start_time_(std::chrono::high_resolution_clock::now()) {
    stats_.total_tasks_submitted.fetch_add(1, std::memory_order_relaxed);
  }

  void operator()() {
    try {
      callable_();

      auto end_time = std::chrono::high_resolution_clock::now();
      auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(
          end_time - start_time_);

      stats_.RecordTaskExecution(priority_, duration);
    } catch (const std::exception& e) {
      stats_.total_tasks_failed.fetch_add(1, std::memory_order_relaxed);
      HandleTaskException(e);
    }
  }

 private:
  Callable callable_;
  TaskPriority priority_;
  TbbTaskStats& stats_;
  std::chrono::high_resolution_clock::time_point start_time_;

  void HandleTaskException(const std::exception& e) {
    // 错误处理逻辑
    // 可以集成到框架的错误恢复系统
  }
};

// TBB原生的高级任务调度器
class TbbAdvancedTaskScheduler {
 public:
  explicit TbbAdvancedTaskScheduler(
      size_t max_threads = tbb::info::default_concurrency())
      : max_threads_(max_threads),
        global_control_(tbb::global_control::max_allowed_parallelism,
                        max_threads),
        main_arena_(static_cast<int>(max_threads)),
        high_priority_arena_(std::max(1, static_cast<int>(max_threads / 4))),
        critical_arena_(1)  // 关键任务使用单独的arena
        ,
        task_stats_() {
    // 启动任务统计更新
    StartStatsUpdater();
  }

  ~TbbAdvancedTaskScheduler() { Shutdown(); }

  // 提交不同优先级的任务
  template <typename Callable>
  void SubmitTask(Callable&& callable,
                  TaskPriority priority = TaskPriority::kNormal) {
    auto wrapped_task =
        TbbTaskWrapper(std::forward<Callable>(callable), priority, task_stats_);

    switch (priority) {
      case TaskPriority::kCritical:
        critical_arena_.execute(
            [wrapped_task = std::move(wrapped_task)]() mutable {
              critical_task_group_.run(std::move(wrapped_task));
            });
        break;

      case TaskPriority::kHigh:
        high_priority_arena_.execute(
            [wrapped_task = std::move(wrapped_task)]() mutable {
              high_priority_task_group_.run(std::move(wrapped_task));
            });
        break;

      case TaskPriority::kNormal:
      case TaskPriority::kLow:
      default:
        main_arena_.execute([wrapped_task = std::move(wrapped_task)]() mutable {
          main_task_group_.run(std::move(wrapped_task));
        });
        break;
    }
  }

  // 提交带返回值的任务
  template <typename Callable>
  auto SubmitTaskWithResult(Callable&& callable,
                            TaskPriority priority = TaskPriority::kNormal)
      -> std::future<decltype(callable())> {
    using ReturnType = decltype(callable());
    auto promise = std::make_shared<std::promise<ReturnType>>();
    auto future = promise->get_future();

    auto task_with_promise = [callable = std::forward<Callable>(callable),
                              promise]() mutable {
      try {
        if constexpr (std::is_void_v<ReturnType>) {
          callable();
          promise->set_value();
        } else {
          auto result = callable();
          promise->set_value(std::move(result));
        }
      } catch (...) {
        promise->set_exception(std::current_exception());
      }
    };

    SubmitTask(std::move(task_with_promise), priority);
    return future;
  }

  // 批量任务提交
  template <typename Iterator, typename Callable>
  void SubmitBatchTasks(Iterator begin, Iterator end, Callable&& callable,
                        TaskPriority priority = TaskPriority::kNormal) {
    auto batch_task = [begin, end,
                       callable = std::forward<Callable>(callable)]() {
      tbb::parallel_for_each(begin, end, callable);
    };

    SubmitTask(std::move(batch_task), priority);
  }

  // 等待所有任务完成
  void WaitForAll() {
    critical_task_group_.wait();
    high_priority_task_group_.wait();
    main_task_group_.wait();
  }

  // 动态调整线程数
  void SetMaxThreads(size_t max_threads) {
    if (max_threads != max_threads_) {
      max_threads_ = max_threads;

      // 重新创建global_control
      global_control_.~global_control();
      new (&global_control_) tbb::global_control(
          tbb::global_control::max_allowed_parallelism, max_threads);

      // 重新创建arenas
      main_arena_.~task_arena();
      new (&main_arena_) tbb::task_arena(static_cast<int>(max_threads));

      high_priority_arena_.~task_arena();
      new (&high_priority_arena_)
          tbb::task_arena(std::max(1, static_cast<int>(max_threads / 4)));
    }
  }

  // 获取任务统计
  TbbTaskStats GetTaskStats() const { return task_stats_; }

  // 获取当前线程使用情况
  struct ThreadUsageInfo {
    size_t max_threads;
    size_t active_threads;
    size_t queued_tasks;
    double cpu_utilization;
  };

  ThreadUsageInfo GetThreadUsage() const {
    return {max_threads_, GetActiveThreadCount(), GetQueuedTaskCount(),
            CalculateCpuUtilization()};
  }

  void Shutdown() {
    running_.store(false, std::memory_order_release);

    // 等待所有任务完成
    WaitForAll();

    // 停止统计更新器
    if (stats_updater_task_group_.is_current_task_group_executing()) {
      stats_updater_task_group_.cancel();
    }
    stats_updater_task_group_.wait();
  }

 private:
  size_t max_threads_;
  tbb::global_control global_control_;

  // 不同优先级的arena和task_group
  tbb::task_arena main_arena_;
  tbb::task_arena high_priority_arena_;
  tbb::task_arena critical_arena_;

  tbb::task_group main_task_group_;
  tbb::task_group high_priority_task_group_;
  tbb::task_group critical_task_group_;

  // 统计信息
  TbbTaskStats task_stats_;
  std::atomic<bool> running_{true};

  // 统计更新器
  tbb::task_group stats_updater_task_group_;

  void StartStatsUpdater() {
    stats_updater_task_group_.run([this]() {
      while (running_.load(std::memory_order_acquire)) {
        UpdateRuntimeStats();
        std::this_thread::sleep_for(std::chrono::seconds(1));
      }
    });
  }

  void UpdateRuntimeStats() {
    // 更新运行时统计信息
    // 可以集成到动态配置系统
  }

  size_t GetActiveThreadCount() const {
    // 获取当前活跃线程数
    // 这里可以通过TBB的内部API或者系统调用获取
    return max_threads_;  // 简化实现
  }

  size_t GetQueuedTaskCount() const {
    // 获取队列中的任务数
    // TBB没有直接提供这个接口，需要自己维护计数器
    return 0;  // 简化实现
  }

  double CalculateCpuUtilization() const {
    // 计算CPU利用率
    return 0.0;  // 简化实现
  }
};
}  // namespace itbb

#endif  // ITBB_CORE_TBB_TASK_SCHEDULER_H_
